import os
import json
import sys
import platform

ascend_toolkit_path = os.environ["ASCEND_HOME_PATH"]

def fix_ascend_toolkit_path(op_path, cross_compile, soc_version):
    with open(f"{op_path}/CMakePresets.json") as fp:
        data = json.load(fp)
    cache = data["configurePresets"][0]["cacheVariables"]
    cache["ASCEND_CANN_PACKAGE_PATH"]["value"] = ascend_toolkit_path
    cache["ENABLE_CROSS_COMPILE"]["value"] = cross_compile and "True" or "False"
    if os.path.exists("/home/ma-user/anaconda3/envs/MindSpore/bin/python"):
        cache["ASCEND_PYTHON_EXECUTABLE"]["value"] = "/home/ma-user/anaconda3/envs/MindSpore/bin/python"
    
    elif os.environ.get("ASCEND_PYTHON_EXECUTABLE"):
        cache["ASCEND_PYTHON_EXECUTABLE"]["value"] = os.environ.get("ASCEND_PYTHON_EXECUTABLE")
    else:
        cache["ASCEND_PYTHON_EXECUTABLE"]["value"] = "/usr/bin/python3"
    cache["ASCEND_COMPUTE_UNIT"]["value"] = soc_version

    if platform.machine() == "x86_64":
        if os.path.exists("/usr/bin/aarch64-linux-gnu-g++"):
            cache["CMAKE_CROSS_PLATFORM_COMPILER"]["value"] = "/usr/bin/aarch64-linux-gnu-g++"
        elif os.path.exists("/usr/bin/aarch64-linux-gnu-gcc-13"):
            cache["CMAKE_CROSS_PLATFORM_COMPILER"]["value"] = "/usr/bin/aarch64-linux-gnu-gcc-13"
    elif platform.machine() == "aarch64":
        cache["CMAKE_CROSS_PLATFORM_COMPILER"]["value"] = "/usr/bin/g++"

    with open(f"{op_path}/CMakePresets.json", "w") as fp:
        json.dump(data, fp=fp, indent="\t")


if __name__ == "__main__":
    op_path = sys.argv[1]
    cross_compile = sys.argv[2] == "True"
    try:
        soc_version = sys.argv[3]
    except:
        soc_version = "ascend310b"
    fix_ascend_toolkit_path(op_path, cross_compile, soc_version)